#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基础模型类
包含通用的数据库操作方法和工具函数
"""

from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime, Text, Boolean, ForeignKey, UniqueConstraint, Index
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql import func
from datetime import datetime
from typing import List, Optional, ClassVar, Dict
import os
import pandas as pd
import json
from loguru import logger

# 创建基类
Base = declarative_base()

# ========================================
# 工具函数
# ========================================

def safe_convert_numeric(value, default=0.0):
    """安全转换数值类型"""
    if pd.isna(value) or value is None:
        return default
    try:
        if isinstance(value, str):
            # 处理可能的无效字符串
            value = value.strip()
            if not value or value == '':
                return default
            # 尝试转换为浮点数
            return float(value)
        elif isinstance(value, (int, float)):
            return float(value)
        else:
            return default
    except (ValueError, TypeError):
        return default

def safe_convert_int(value, default=0):
    """安全转换整数类型"""
    if pd.isna(value) or value is None:
        return default
    try:
        if isinstance(value, str):
            value = value.strip()
            if not value or value == '':
                return default
            return int(float(value))  # 先转浮点数再转整数
        elif isinstance(value, (int, float)):
            return int(value)
        else:
            return default
    except (ValueError, TypeError):
        return default

def safe_get_field_value(row, field_name, default=None):
    """安全获取字段值，如果字段不存在则返回默认值"""
    if field_name in row:
        value = row[field_name]
        if pd.isna(value) or value is None:
            return default
        return value
    return default

def standardize_date_format(date_value):
    """
    统一日期格式为 YYYY-MM-DD
    
    Args:
        date_value: 日期值，可能是字符串、数字或datetime对象
        
    Returns:
        str: 标准化的日期字符串 YYYY-MM-DD，如果转换失败返回None
    """
    if date_value is None:
        return None
    
    try:
        # 如果是字符串，先尝试解析
        if isinstance(date_value, str):
            date_str = str(date_value).strip()
            if not date_str or date_str == '0' or date_str == '':
                return None
            
            # 如果已经是 YYYY-MM-DD 格式，验证日期有效性
            if len(date_str) == 10 and date_str[4] == '-' and date_str[7] == '-':
                # 验证日期有效性
                import pandas as pd
                parsed_date = pd.to_datetime(date_str, errors='coerce')
                if pd.notna(parsed_date):
                    return date_str
                return None
            
            # 如果是 YYYYMMDD 格式，转换为 YYYY-MM-DD
            if len(date_str) == 8 and date_str.isdigit():
                year = int(date_str[:4])
                month = int(date_str[4:6])
                day = int(date_str[6:8])
                # 验证日期有效性
                import pandas as pd
                try:
                    parsed_date = pd.to_datetime(f"{year:04d}-{month:02d}-{day:02d}", errors='coerce')
                    if pd.notna(parsed_date):
                        return f"{year:04d}-{month:02d}-{day:02d}"
                except:
                    pass
                return None
            
            # 尝试其他格式解析
            import pandas as pd
            parsed_date = pd.to_datetime(date_str, errors='coerce')
            if pd.notna(parsed_date):
                return parsed_date.strftime('%Y-%m-%d')
        
        # 如果是数字，转换为字符串后处理
        elif isinstance(date_value, (int, float)):
            if date_value <= 0:
                return None
            date_str = str(int(date_value))
            if len(date_str) == 8:
                year = int(date_str[:4])
                month = int(date_str[4:6])
                day = int(date_str[6:8])
                # 验证日期有效性
                import pandas as pd
                try:
                    parsed_date = pd.to_datetime(f"{year:04d}-{month:02d}-{day:02d}", errors='coerce')
                    if pd.notna(parsed_date):
                        return f"{year:04d}-{month:02d}-{day:02d}"
                except:
                    pass
                return None
        
        # 如果是datetime对象
        elif hasattr(date_value, 'strftime'):
            return date_value.strftime('%Y-%m-%d')
        
        return None
        
    except Exception as e:
        # 日期格式转换失败: {date_value}, 错误: {str(e)}
        return None

# ========================================
# 数据库管理器
# ========================================

class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self, stock_selection_db_path: str = "stock_selection.db", 
                 stock_data_db_path: str = "stock_data.db"):
        """初始化数据库管理器"""
        self.stock_selection_db_path = stock_selection_db_path
        self.stock_data_db_path = stock_data_db_path
        
        # 创建数据库引擎
        self.stock_selection_engine = create_engine(f'sqlite:///{stock_selection_db_path}')
        self.stock_data_engine = create_engine(f'sqlite:///{stock_data_db_path}')
        
        # 创建会话工厂
        self.StockSelectionSession = sessionmaker(bind=self.stock_selection_engine)
        self.StockDataSession = sessionmaker(bind=self.stock_data_engine)
    
    def init_databases(self):
        """初始化所有数据库表"""
        from .HSnapshot import HSnapshot
        from .Block import Block
        from .BlockInfo import BlockInfo
        from .FBoard import FBoard
        from .SelectResult import SelectResult
        from .ExecutionCache import ExecutionCache
        from .NewsData import NewsData
        from .StockInfo import StockInfo
        from .LabelInfo import LabelInfo
        
        # 创建选股结果数据库表
        Base.metadata.create_all(self.stock_selection_engine, 
                                tables=[SelectResult.__table__, ExecutionCache.__table__])
        
        # 创建股票数据数据库表
        Base.metadata.create_all(self.stock_data_engine,
                                tables=[HSnapshot.__table__, Block.__table__, 
                                       BlockInfo.__table__, FBoard.__table__,
                                       NewsData.__table__, StockInfo.__table__, LabelInfo.__table__])
        
        # 数据库表初始化完成
    
    def get_stock_selection_session(self):
        """获取选股结果数据库会话"""
        return self.StockSelectionSession()
    
    def get_stock_data_session(self):
        """获取股票数据数据库会话"""
        return self.StockDataSession()
    
    def close_all_sessions(self):
        """关闭所有数据库连接"""
        self.stock_selection_engine.dispose()
        self.stock_data_engine.dispose()

# ========================================
# 基础模型类
# ========================================

class BaseModel(Base):
    """基础模型类，提供通用的数据库操作方法"""
    __abstract__ = True
    
    # 类变量：数据库管理器
    _db_manager: ClassVar[Optional[DatabaseManager]] = None
    
    @classmethod
    def set_db_manager(cls, db_manager: DatabaseManager):
        """设置数据库管理器"""
        cls._db_manager = db_manager
    
    @classmethod
    def get_session(cls):
        """获取数据库会话（子类需要重写）"""
        raise NotImplementedError("子类必须实现get_session方法")
    
    @classmethod
    def save_batch(cls, data_list: List[dict]) -> bool:
        """批量保存数据（通用方法）"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            for data in data_list:
                # 检查是否已存在相同记录（子类需要实现）
                existing = cls._get_existing_record(session, data)
                
                if existing:
                    # 更新现有记录
                    cls._update_existing_record(existing, data)
                else:
                    # 创建新记录
                    new_record = cls(**data)
                    session.add(new_record)
            
            session.commit()
            session.close()
            return True
            
        except Exception as e:
            # 批量保存数据失败: {str(e)}
            session.rollback()
            session.close()
            return False
    
    @classmethod
    def _get_existing_record(cls, session, data):
        """获取现有记录（子类需要重写）"""
        raise NotImplementedError("子类必须实现_get_existing_record方法")
    
    @classmethod
    def _update_existing_record(cls, existing, data):
        """更新现有记录（子类需要重写）"""
        raise NotImplementedError("子类必须实现_update_existing_record方法")
    
    @classmethod
    def to_dataframe(cls, **kwargs) -> pd.DataFrame:
        """转换为DataFrame（通用方法）"""
        if not cls._db_manager:
            raise ValueError("数据库管理器未设置，请先调用 set_db_manager()")
        
        try:
            session = cls.get_session()
            
            # 构建查询（子类可以重写）
            query = cls._build_query(session, **kwargs)
            results = query.all()
            
            # 转换为DataFrame
            data_list = []
            for result in results:
                data_list.append(cls._to_dict(result))
            
            session.close()
            
            df = pd.DataFrame(data_list)
            # 从数据库加载了 {len(df)} 条数据
            return df
            
        except Exception as e:
            # 转换为DataFrame失败: {str(e)}
            session.close()
            return pd.DataFrame()
    
    @classmethod
    def _build_query(cls, session, **kwargs):
        """构建查询（子类可以重写）"""
        return session.query(cls)
    
    @classmethod
    def _to_dict(cls, result):
        """将结果转换为字典（子类需要重写）"""
        raise NotImplementedError("子类必须实现_to_dict方法")
    
    def to_dict(self) -> dict:
        """转换为字典（通用方法）"""
        return self._to_dict(self)
